import * as THREE from "three"

const gravity = new THREE.Vector3(0, -0.005, 0);
const friction = 0.998;
const textureSize = 128.0;

const rangeNum = (max = 1, min = 0) => Math.floor(Math.random() * (max + 1 - min)) + min;

const getOffsetXYZ = i => {
    const offset = 3;
    const index = i * offset;
    const x = index;
    const y = index + 1;
    const z = index + 2;
    return { x, y, z };
};

const getOffsetRGBA = i => {
    const offset = 4;
    const index = i * offset;
    const r = index;
    const g = index + 1;
    const b = index + 2;
    const a = index + 3;
    return { r, g, b, a };
};

export class ParticleMesh extends THREE.Points {
    /**
     * @param {number} num 
     * @param {THREE.Vector3[]} _velocities 
     * @param {"trail"|"seed"} type 
     */
    constructor(num, _velocities, type) {

        super()

        this.particleNum = num;
        this.timerStartFading = 10;

        // geometry

        const bufferGeometry = this.geometry

        const vertices = new Float32Array(num * 3);
        const velocities = new Float32Array(num * 3);
        const colors = new Float32Array(num * 4);
        const adjustSizes = new Float32Array(num);
        const masses = new Float32Array(num);

        const colorType = Math.random() > 0.3 ? 'single' : 'multiple';
        const singleColor = rangeNum(1, 0.3);
        const multipleColor = () => rangeNum(1, 0.01);
        let rgbType;
        const rgbTypeDice = Math.random();
        if (rgbTypeDice > 0.66) {
            rgbType = 'red';
        } else if (rgbTypeDice > 0.33) {
            rgbType = 'green';
        } else {
            rgbType = 'blue';
        }
        for (let i = 0; i < num; i++) {
            const pos = new THREE.Vector3(0, 0, 0);

            const { x, y, z } = getOffsetXYZ(i)
            const { r, g, b, a } = getOffsetRGBA(i)

            vertices[x] = pos.x
            vertices[y] = pos.y
            vertices[z] = pos.z

            const velocity = _velocities[i]
            velocities[x] = velocity.x
            velocities[y] = velocity.y
            velocities[z] = velocity.z

            if (type === 'seed') {

                // 烟花种子
                let size = Math.pow(velocity.y, 2) * 0.04;
                if (i === 0) size *= 1.1;
                adjustSizes[i] = size;
                masses[i] = size * 0.017
                colors[r] = 1.0
                colors[g] = 1.0
                colors[b] = 1.0
                colors[a] = 1.0
            } else {

                // 烟花爆炸后的花
                const size = rangeNum(30, 1) * 0.01;
                adjustSizes[i] = size;
                masses[i] = size * 0.017
                if (colorType === 'multiple') {

                    colors[r] = multipleColor()
                    colors[g] = multipleColor()
                    colors[b] = multipleColor()
                    colors[a] = 1.0

                } else {
                    switch (rgbType) {
                        case 'red':
                            colors[r] = singleColor
                            colors[g] = 0.1
                            colors[b] = 0.1
                            colors[a] = 1.0
                            break;
                        case 'green':
                            colors[r] = 0.1
                            colors[g] = singleColor
                            colors[b] = 0.1
                            colors[a] = 1.0
                            break;
                        case 'blue':
                            colors[r] = 0.1
                            colors[g] = 0.1
                            colors[b] = singleColor
                            colors[a] = 1.0
                            break;
                        default:
                            colors[r] = singleColor
                            colors[g] = 0.1
                            colors[b] = 0.1
                            colors[a] = 1.0
                    }
                }
            }
        }
        bufferGeometry.setAttribute('position', new THREE.BufferAttribute(vertices, 3));
        bufferGeometry.setAttribute('velocity', new THREE.BufferAttribute(velocities, 3));
        bufferGeometry.setAttribute('color', new THREE.BufferAttribute(colors, 4));
        bufferGeometry.setAttribute('adjustSize', new THREE.BufferAttribute(adjustSizes, 1));
        bufferGeometry.setAttribute('mass', new THREE.BufferAttribute(masses, 1));
        // material
        this.material = new THREE.RawShaderMaterial({
            uniforms: {
                size: {
                    value: textureSize
                },
            },
            transparent: true,
            // Display of "blending: THREE.AdditiveBlending" does not work properly if "depthWrite" property is set to true.
            // Therefore, it is necessary to make it false in the case of making the image transparent by blending.
            depthWrite: false,
            blending: THREE.AdditiveBlending,
            vertexShader: document.getElementById('vs').textContent,
            fragmentShader: document.getElementById('fs').textContent
        });
    }
    update(gravity) {

        if (this.timerStartFading > 0) this.timerStartFading -= 0.3;

        const { position, velocity, color, mass } = this.geometry.attributes;

        const decrementRandom = () => (Math.random() > 0.5 ? 0.98 : 0.96);

        const decrementByVel = v => (Math.random() > 0.5 ? 0 : (1 - v) * 0.1);

        for (let i = 0; i < this.particleNum; i++) {

            const { x, y, z } = getOffsetXYZ(i);
            velocity.array[y] += gravity.y - mass.array[i];
            velocity.array[x] *= friction;
            velocity.array[z] *= friction;
            velocity.array[y] *= friction;
            position.array[x] += velocity.array[x];
            position.array[y] += velocity.array[y];
            position.array[z] += velocity.array[z];
            const { a } = getOffsetRGBA(i);
            if (this.timerStartFading <= 0) {
                color.array[a] *= decrementRandom() - decrementByVel(color.array[a]);
                if (color.array[a] < 0.001) color.array[a] = 0;
            }
        }
        position.needsUpdate = true;
        velocity.needsUpdate = true;
        color.needsUpdate = true;
    }
    dispose() {
        this.geometry.dispose();
        this.material.dispose();
    }
}

export class ParticleSeedMesh extends ParticleMesh {
    constructor(num, _velocities) {
        super(num, _velocities, 'seed');
    }
    update(gravity) {
        const { position, velocity, color, mass } = this.geometry.attributes;
        const decrementRandom = () => (Math.random() > 0.3 ? 0.99 : 0.96);
        const decrementByVel = v => (Math.random() > 0.3 ? 0 : (1 - v) * 0.1);
        const shake = () => Math.random() > 0.5 ? 0.05 : -0.05;
        const dice = () => Math.random() > 0.1;
        const _f = friction * 0.98;
        for (let i = 0; i < this.particleNum; i++) {
            const { x, y, z } = getOffsetXYZ(i);
            velocity.array[y] += gravity.y - mass.array[i];
            velocity.array[x] *= _f;
            velocity.array[z] *= _f;
            velocity.array[y] *= _f;
            position.array[x] += velocity.array[x];
            position.array[y] += velocity.array[y];
            position.array[z] += velocity.array[z];
            if (dice()) position.array[x] += shake();
            if (dice()) position.array[z] += shake();
            const { a } = getOffsetRGBA(i);
            color.array[a] *= decrementRandom() - decrementByVel(color.array[a]);
            if (color.array[a] < 0.001) color.array[a] = 0;
        }
        position.needsUpdate = true;
        velocity.needsUpdate = true;
        color.needsUpdate = true;
    }
}

export class ParticleTailMesh extends ParticleMesh {
    constructor(num, _velocities) {
        super(num, _velocities, 'trail');
    }
    update(gravity) {
        const { position, velocity, color, mass } = this.geometry.attributes;
        const decrementRandom = () => (Math.random() > 0.3 ? 0.98 : 0.95);
        const shake = () => (Math.random() > 0.5 ? 0.05 : -0.05);
        const dice = () => Math.random() > 0.2;
        for (let i = 0; i < this.particleNum; i++) {
            const { x, y, z } = getOffsetXYZ(i);
            velocity.array[y] += gravity.y - mass.array[i];
            velocity.array[x] *= friction;
            velocity.array[z] *= friction;
            velocity.array[y] *= friction;
            position.array[x] += velocity.array[x];
            position.array[y] += velocity.array[y];
            position.array[z] += velocity.array[z];
            if (dice()) position.array[x] += shake();
            if (dice()) position.array[z] += shake();
            const { a } = getOffsetRGBA(i);
            color.array[a] *= decrementRandom();
            if (color.array[a] < 0.001) color.array[a] = 0;
        }
        position.needsUpdate = true;
        velocity.needsUpdate = true;
        color.needsUpdate = true;
    }
}

export class BasicFIreWorks extends THREE.Group {
    constructor() {
        super()
        this.isExplode = false;
        const max = 400;
        const min = 150;
        this.petalsNum = rangeNum(max, min);
        this.life = 150;
        this.seed = this.getSeed();

        this.flower = null

        this.add(this.seed)
        this.flowerSizeRate = THREE.MathUtils.mapLinear(this.petalsNum, min, max, 0.4, 0.7);

    }
    getSeed() {
        const num = 40;
        const _velocities = [];
        for (let i = 0; i < num; i++) {
            const vx = 0;
            const vy = i === 0 ? Math.random() * 2.5 + 0.9 : Math.random() * 2.0 + 0.4;
            const vz = 0;
            _velocities.push(new THREE.Vector3(vx, vy, vz));
        }
        const pm = new ParticleSeedMesh(num, _velocities);
        const x = Math.random() * 80 - 40;
        const y = -50;
        const z = Math.random() * 80 - 40;
        pm.position.set(x, y, z);
        return pm;
    }
    explode(pos) {
        this.isExplode = true;
        this.flower = this.getFlower(pos);
        this.add(this.flower);
        this.remove(this.seed);
        this.seed.dispose();
    }
    getFlower(pos) {
        const num = this.petalsNum;
        const _velocities = [];
        let radius;
        const dice = Math.random();

        if (dice > 0.5) {
            for (let i = 0; i < num; i++) {
                radius = rangeNum(120, 60) * 0.01;
                const theta = THREE.MathUtils.degToRad(Math.random() * 180);
                const phi = THREE.MathUtils.degToRad(Math.random() * 360);
                const vx = Math.sin(theta) * Math.cos(phi) * radius;
                const vy = Math.sin(theta) * Math.sin(phi) * radius;
                const vz = Math.cos(theta) * radius;
                const vel = new THREE.Vector3(vx, vy, vz);
                vel.multiplyScalar(this.flowerSizeRate);
                _velocities.push(vel);
            }
        } else {
            const zStep = 180 / num;
            const trad = (360 * (Math.random() * 20 + 1)) / num;
            const xStep = trad;
            const yStep = trad;
            radius = rangeNum(120, 60) * 0.01;
            for (let i = 0; i < num; i++) {
                const sphereRate = Math.sin(THREE.MathUtils.degToRad(zStep * i));
                const vz = Math.cos(THREE.MathUtils.degToRad(zStep * i)) * radius;
                const vx = Math.cos(THREE.MathUtils.degToRad(xStep * i)) * sphereRate * radius;
                const vy = Math.sin(THREE.MathUtils.degToRad(yStep * i)) * sphereRate * radius;
                const vel = new THREE.Vector3(vx, vy, vz);
                vel.multiplyScalar(this.flowerSizeRate);
                _velocities.push(vel);
            }
        }

        const particleMesh = new ParticleMesh(num, _velocities);
        particleMesh.position.set(pos.x, pos.y, pos.z);
        return particleMesh;
    }
    update(gravity) {
        if (!this.isExplode) {
            this.drawTail();
        } else {
            this.flower.update(gravity);
            if (this.life > 0) this.life -= 1;
        }
    }
    drawTail() {
        this.seed.update(gravity);
        const { position, velocity } = this.seed.geometry.attributes;
        let count = 0;
        let isComplete = true;
        // Check if the y-axis speed is down for all particles
        for (let i = 0, l = velocity.array.length; i < l; i++) {
            const v = velocity.array[i];
            const index = i % 3;
            if (index === 1 && v > 0) {
                count++;
            }
        }

        isComplete = count == 0;
        if (!isComplete) return;
        const { x, y, z } = this.seed.position;
        const flowerPos = new THREE.Vector3(x, y, z);
        let highestPos = 0;
        let offsetPos;
        for (let i = 0, l = position.array.length; i < l; i++) {
            const p = position.array[i];
            const index = i % 3;
            if (index === 1 && p > highestPos) {
                highestPos = p;
                offsetPos = new THREE.Vector3(position.array[i - 1], p, position.array[i + 2]);
            }
        }
        flowerPos.add(offsetPos);
        this.explode(flowerPos);
    }
}

export class RichFIreWorks extends BasicFIreWorks {
    constructor() {
        super();
        const max = 150;
        const min = 100;
        this.petalsNum = rangeNum(max, min);
        this.flowerSizeRate = THREE.MathUtils.mapLinear(this.petalsNum, min, max, 0.4, 0.7);
        this.tailMeshGroup = new THREE.Group();
        this.tails = [];
    }
    explode(pos) {
        this.isExplode = true;
        this.flower = this.getFlower(pos);
        this.tails = this.getTail();
        this.add(this.flower);
        this.add(this.tailMeshGroup);
    }
    getTail() {
        const tails = [];
        const num = 20;
        const { color: petalColor } = this.flower.geometry.attributes;

        for (let i = 0; i < this.petalsNum; i++) {
            const _velocities = [];
            for (let j = 0; j < num; j++) {
                const vx = 0;
                const vy = 0;
                const vz = 0;
                _velocities.push(new THREE.Vector3(vx, vy, vz));
            }
            const tail = new ParticleTailMesh(num, _velocities);

            const { r, g, b, a } = getOffsetRGBA(i);

            const petalR = petalColor.array[r];
            const petalG = petalColor.array[g];
            const petalB = petalColor.array[b];
            const petalA = petalColor.array[a];

            const { position, color } = tail.geometry.attributes;

            for (let k = 0; k < position.count; k++) {
                const { r, g, b, a } = getOffsetRGBA(k);
                color.array[r] = petalR;
                color.array[g] = petalG;
                color.array[b] = petalB;
                color.array[a] = petalA;
            }

            const { x, y, z } = this.flower.position;
            tail.position.set(x, y, z);
            tails.push(tail);
            this.tailMeshGroup.add(tail);
        }
        return tails;
    }
    update(gravity) {
        if (!this.isExplode) {
            this.drawTail();
        } else {
            this.flower.update(gravity);

            const { position: flowerGeometry } = this.flower.geometry.attributes;

            for (let i = 0, l = this.tails.length; i < l; i++) {
                const tail = this.tails[i];
                tail.update(gravity);
                const flowerPos = new THREE.Vector3().fromBufferAttribute(flowerGeometry, i);
                const { position, velocity } = tail.geometry.attributes;
                for (let k = 0; k < position.count; k++) {
                    const { x, y, z } = getOffsetXYZ(k);
                    const desiredVelocity = new THREE.Vector3();
                    const tailPos = new THREE.Vector3(position.array[x], position.array[y], position.array[z]);
                    const tailVel = new THREE.Vector3(velocity.array[x], velocity.array[y], velocity.array[z]);
                    desiredVelocity.subVectors(flowerPos, tailPos);
                    const steer = desiredVelocity.sub(tailVel);
                    steer.normalize();
                    steer.multiplyScalar(Math.random() * 0.0003 * this.life);
                    velocity.array[x] += steer.x;
                    velocity.array[y] += steer.y;
                    velocity.array[z] += steer.z;
                }
                velocity.needsUpdate = true;
            }

            if (this.life > 0) this.life -= 1.2;
        }
    }
}